# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import math
import sys

import tabulate

from coin.base.param_util import to_list
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.okex_futures.types import OkexFuturesOrderSide
from coin.exchange.okex_futures.kr_rest.currency import OkexCurrency
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFutureProduct
from coin.strategy.mm.calculator import MovingSum
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.strategy.mm.bitmex.simple_mm_executor import BitmexSimpleMMExecutor
from coin.strategy.mm.base.okex_futures_bitmex_strategy_base import StrategyBase


def ceil_price(price, tick_size):
  return math.ceil(price / tick_size) * tick_size


def floor_price(price, tick_size):
  return math.floor(price / tick_size) * tick_size


def norm_float(value, precision):
  return float(('%%.%df' % precision) % value)


def calculate_unrealized_pnl(pos, avg_price, ask0_price, bid0_price):
  if pos == 0:
    return 0
  mid_price = (ask0_price + bid0_price) / 2.
  return pos * (1. / avg_price - 1. / mid_price)


class SoftArbStrategy(StrategyBase):
  def __init__(self, strategy_params, dry_run=True):
    super().__init__(product_1=strategy_params['product_agg'],
                     product_2=strategy_params['product_passive'],
                     balance_currency=strategy_params.get('balance_currency', None),
                     logger=logging.getLogger('SoftArbStrategy'))

    self._strategy_params = strategy_params
    self._dry_run = dry_run

    self._edge = strategy_params['edge_bp'] / 10000.
    self._basis_ma_window = strategy_params['basis_ma_window_min'] * (60 * 10**9)

    self._true_book_2 = strategy_params['passive_true_book']
    self._lot_size_2 = strategy_params['lot_size']
    self._tick_size_2 = strategy_params['tick_size']

    self._agg_edge = strategy_params['agg_edge_bp'] / 10000.
    self._true_book_1 = strategy_params['agg_true_book']
    self._lot_size_min_1 = strategy_params['agg_lot_size_min']
    self._lot_size_max_1 = strategy_params['agg_lot_size_max']
    self._pos_multiplier_1 = strategy_params['agg_pos_multiplier']

    self._passive_executor = None
    self._ask_basis_ma = MovingSum(self._basis_ma_window)
    self._bid_basis_ma = MovingSum(self._basis_ma_window)

  def on_base_init(self, ioloop, queue, flow_sub):
    super().on_base_init(ioloop, queue, flow_sub)
    self.start_print_feature_status()

  def init_executor(self):
    # Intialize passive executor
    self._logger.info('Intializing passive executor...')
    passive_executor_params = {
        'product': self._strategy_params['product_passive'], 'order_update_period': 3.0
    }
    for field in ['lot_size', 'stickiness', 'min_pos', 'max_pos']:
      passive_executor_params[field] = self._strategy_params[field]
    self._passive_executor = BitmexSimpleMMExecutor(passive_executor_params,
                                                    self._og_2,
                                                    self._logger)

    # Register fill callback
    self._og_1.register_fill_callback(self.on_agg_fill)
    self._og_2.register_fill_callback(self.on_passive_fill)

  def start_print_feature_status(self):
    try:
      ask_basis_ma = (self._ask_basis_ma.average if len(self._ask_basis_ma) > 0 else float('nan'))
      bid_basis_ma = (self._bid_basis_ma.average if len(self._bid_basis_ma) > 0 else float('nan'))
      self._logger.info(
          '-----\n- Feature status\n'
          '-   ask_basis_ma=%.4f, len=%d\n'
          '-   bid_basis_ma=%.4f, len=%d',
          ask_basis_ma,
          len(self._ask_basis_ma),
          bid_basis_ma,
          len(self._bid_basis_ma))

    finally:
      self._ioloop.add_timeout(datetime.timedelta(seconds=30), self.start_print_feature_status)

  def _print_balance(self):  # Override the original version.
    if self._book_1 is None or self._book_2 is None:
      return

    balance_1 = self._og_1.get_balance('BTC')  # Not includes unrealized pnl

    pos_from_og_1 = self._og_1.get_position(self._product_1)
    if pos_from_og_1 is None:
      unrealized_pnl_1 = 0
      pos_1 = 0
    else:
      long_pos_1 = pos_from_og_1.long * self._pos_multiplier_1
      short_pos_1 = -pos_from_og_1.short * self._pos_multiplier_1
      long_unrealized_pnl_1 = calculate_unrealized_pnl(long_pos_1,
                                                       pos_from_og_1.long_avg_price,
                                                       self._book_1.ask0().price,
                                                       self._book_1.bid0().price)
      short_unrealized_pnl_1 = calculate_unrealized_pnl(short_pos_1,
                                                        pos_from_og_1.short_avg_price,
                                                        self._book_1.ask0().price,
                                                        self._book_1.bid0().price)
      unrealized_pnl_1 = long_unrealized_pnl_1 + short_unrealized_pnl_1
      pos_1 = long_pos_1 + short_pos_1

    balance_1 += unrealized_pnl_1

    balance_2 = self._og_2.get_balance('BTC')  # Includes unrealized pnl
    pos_2 = self._og_2.get_position('BTC-USD.PERPETUAL')

    table = [['*', 'BTC', balance_1 + balance_2, balance_1, balance_2],
             ['*', 'USD', pos_1 + pos_2, pos_1, pos_2]]

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table, headers=['*', 'Currency', 'Total', 'OKEX Futures', 'Bitmex'], floatfmt=".6f")))

  def on_book(self):
    if self._book_1 is None or self._book_2 is None:
      return

    # Check book freshness
    book_ts = max(self._book_1.timestamp, self._book_2.timestamp)
    min_ts = min(self._book_1.timestamp, self._book_2.timestamp)
    if book_ts - min_ts > (2 * 10**9):
      return

    # Check book sanity
    if (self._book_1.ask0().price <= self._book_1.bid0().price
        or self._book_2.ask0().price <= self._book_2.bid0().price):
      return

    askt_1 = self._book_1.get_notional_asks_by_qty(self._true_book_1
                                                   / self._pos_multiplier_1)[1][-1][0]
    bidt_1 = self._book_1.get_notional_bids_by_qty(self._true_book_1
                                                   / self._pos_multiplier_1)[1][-1][0]

    ask0_2 = self._book_2.ask0().price
    bid0_2 = self._book_2.bid0().price
    askt_2 = self._book_2.get_notional_asks_by_qty(self._true_book_2)[1][-1][0]
    bidt_2 = self._book_2.get_notional_bids_by_qty(self._true_book_2)[1][-1][0]

    ask_diff = askt_1 - askt_2
    bid_diff = bidt_1 - bidt_2

    self._ask_basis_ma.update(book_ts, ask_diff)
    self._bid_basis_ma.update(book_ts, bid_diff)
    if len(self._ask_basis_ma) < 1000 or len(self._bid_basis_ma) < 1000:
      return

    askt_1_proj_2 = askt_1 - float(self._ask_basis_ma.average)
    bidt_1_proj_2 = bidt_1 - float(self._bid_basis_ma.average)

    askp_2 = max(askt_1_proj_2 * (1. + self._edge),
                 ask0_2 - self._tick_size_2,
                 bid0_2 + self._tick_size_2)
    bidp_2 = min(bidt_1_proj_2 * (1. - self._edge),
                 bid0_2 + self._tick_size_2,
                 ask0_2 - self._tick_size_2)
    askp_2 = ceil_price(askp_2, self._tick_size_2)
    bidp_2 = floor_price(bidp_2, self._tick_size_2)

    if not self.is_og_ready() or self._dry_run:
      return

    fired = self._passive_executor.manage_mm_orders(askp_2, bidp_2)

    if fired:
      # If there is no new passive order fired, do not cancel aggression orders.
      self.cancel_all_aggression_orders()
      self._logger.info('[2] SELL %f @ %.1f (%.1f), BUY %f @ %.1f (%.1f)',
                        self._lot_size_2,
                        askp_2,
                        askp_2 - ask0_2,
                        self._lot_size_2,
                        bidp_2,
                        bid0_2 - bidp_2)

  def on_passive_fill(self, fill_price, fill_qty, order):
    if (order.product != self._product_2.native_symbol or self._dry_run):
      return

    if not self.is_og_ready() or self._dry_run:
      self._logger.error('OG is not ready. Skip aggression')
      return

    if order.order_side == OrderSide.SELL:
      self.try_aggression_buy(fill_price, fill_qty)
    elif order.order_side == OrderSide.BUY:
      self.try_aggression_sell(fill_price, fill_qty)
    else:
      raise ValueError('Invalid side: %s' % str(order.order_side))

  def on_agg_fill(self, avg_price, last_fill_qty, order):
    self._logger.info('[1] FILL %s %f @ %.4f', str(order.order_side), last_fill_qty, avg_price)

  def try_aggression_buy(self, fill_price, fill_qty):
    if self._dry_run:
      return

    askt_1 = self._book_1.get_notional_asks_by_qty(self._true_book_1
                                                   / self._pos_multiplier_1)[1][-1][0]
    askt_1_proj_2 = askt_1 - float(self._ask_basis_ma.average)
    if (fill_price / askt_1_proj_2) < (1. + self._agg_edge):
      self._logger.error(
          '[1] SELL-BUY: Not enough edge: '
          'fill=%.1f, askt_1=%.4f, askt_1_proj_2=%.4f, edge=%.2fbp',
          fill_price,
          askt_1,
          askt_1_proj_2, (fill_price / askt_1_proj_2 - 1.) * 10000.0)
      return

    pos_from_og_1 = self._og_1.get_position(self._product_1)
    if pos_from_og_1 is None:
      long_pos_1 = 0
      short_pos_1 = 0
      pos_1 = 0
    else:
      long_pos_1 = pos_from_og_1.long * self._pos_multiplier_1
      short_pos_1 = -pos_from_og_1.short * self._pos_multiplier_1
      pos_1 = long_pos_1 + short_pos_1

    buy_unacked_qty = 0
    for order in self._og_1.get_working_order():
      if (order.product == self._product_1 and not order.accepted
          and order.order_side == OkexFuturesOrderSide.BUY_OPEN
          or order.order_side == OkexFuturesOrderSide.BUY_CLOSE):
        buy_unacked_qty += order.qty
    pos_1 += buy_unacked_qty * self._pos_multiplier_1

    pos_2 = self._og_2.get_position('BTC-USD.PERPETUAL')
    pos_2 -= fill_qty

    diff = 0 - round((pos_1 + pos_2) / self._pos_multiplier_1) * self._pos_multiplier_1
    if diff < self._lot_size_min_1:
      self._logger.info('[1] NO BUY: qty=%.4f < lot_size_min=%.4f', diff, self._lot_size_min_1)
      return

    buy_qty = min(diff, self._lot_size_max_1)

    order_side = OkexFuturesOrderSide.BUY_OPEN
    if -short_pos_1 >= 2 * self._lot_size_2 and -short_pos_1 >= buy_qty:
      order_side = OkexFuturesOrderSide.BUY_CLOSE

    self._og_1.submit(self._product_1,
                      askt_1,
                      buy_qty / self._pos_multiplier_1,
                      order_side,
                      OrderType.LIMIT)
    self._logger.info('[1] BUY  %.4f @ %.4f', buy_qty, askt_1)

  def try_aggression_sell(self, fill_price, fill_qty):
    if self._dry_run:
      return

    bidt_1 = self._book_1.get_notional_bids_by_qty(self._true_book_1
                                                   / self._pos_multiplier_1)[1][-1][0]
    bidt_1_proj_2 = bidt_1 - float(self._bid_basis_ma.average)
    if (bidt_1_proj_2 / fill_price) < (1. + self._agg_edge):
      self._logger.error(
          '[1] BUY-SELL: Not enough edge: '
          'fill=%.1f, bidt_1=%.4f, bidt_1_proj_2=%.4f, edge=%.2fbp',
          fill_price,
          bidt_1,
          bidt_1_proj_2, (bidt_1_proj_2 / fill_price - 1.) * 10000.0)
      return

    pos_from_og_1 = self._og_1.get_position(self._product_1)
    if pos_from_og_1 is None:
      long_pos_1 = 0
      short_pos_1 = 0
      pos_1 = 0
    else:
      long_pos_1 = pos_from_og_1.long * self._pos_multiplier_1
      short_pos_1 = -pos_from_og_1.short * self._pos_multiplier_1
      pos_1 = long_pos_1 + short_pos_1

    sell_unacked_qty = 0
    for order in self._og_1.get_working_order():
      if (order.product == self._product_1 and not order.accepted
          and order.order_side == OkexFuturesOrderSide.SELL_OPEN
          or order.order_side == OkexFuturesOrderSide.SELL_CLOSE):
        sell_unacked_qty += order.qty
    pos_1 -= sell_unacked_qty * self._pos_multiplier_1

    pos_2 = self._og_2.get_position('BTC-USD.PERPETUAL')
    pos_2 += fill_qty

    diff = 0 - round((pos_1 + pos_2) / self._pos_multiplier_1) * self._pos_multiplier_1
    if -diff < self._lot_size_min_1:
      self._logger.info('[1] NO SELL: qty=%.4f < lot_size_min=%.4f', -diff, self._lot_size_min_1)
      return

    sell_qty = min(-diff, self._lot_size_max_1)

    order_side = OkexFuturesOrderSide.SELL_OPEN
    if long_pos_1 >= 2 * self._lot_size_2 and long_pos_1 >= sell_qty:
      order_side = OkexFuturesOrderSide.SELL_CLOSE

    self._og_1.submit(self._product_1,
                      bidt_1,
                      sell_qty / self._pos_multiplier_1,
                      order_side,
                      OrderType.LIMIT)
    self._logger.info('[1] SELL %.4f @ %.4f', sell_qty, bidt_1)

  def cancel_all_aggression_orders(self):
    if self._dry_run:
      return
    orders_1 = [order for order in self._og_1.get_working_order()]
    for order in orders_1:
      if order.accepted and order.product == self._product_1:
        self._og_1.cancel(order.order_id)


def get_param_set(param_set_name, init_og_immediately=False):
  assert param_set_name in ['quarter', 'this_week']

  # Strategy parameter
  product_passive = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')

  if param_set_name == 'quarter':
    product_agg = OkexFutureProduct(OkexCurrency.FromStrCurrency('BTC'),
                                    OkexCurrency.FromStrCurrency('USD'),
                                    'quarter')
    lot_size = 7000.  # USD
    stack = 20
    edge = 15.
    okex_futures_og_config = 'okex_jshin1'
    bitmex_og_config = 'bitmex_jshin3'

  elif param_set_name == 'this_week':
    product_agg = OkexFutureProduct(OkexCurrency.FromStrCurrency('BTC'),
                                    OkexCurrency.FromStrCurrency('USD'),
                                    'this_week')
    lot_size = 3000.  # USD
    stack = 20
    edge = 15.
    okex_futures_og_config = 'okex_jaewon0'
    bitmex_og_config = 'bitmex_jshin2'

  else:
    raise ValueError()

  strategy_params = {
      'edge_bp': edge,
      'basis_ma_window_min': 30,

  # Passive-posting parameters
      'product_passive': product_passive,
      'passive_true_book': 10000.,
      'lot_size': lot_size,
      'stickiness': 0.,
      'min_pos': -lot_size * stack,
      'max_pos': lot_size * stack,
      'tick_size': 0.5,

  # Aggression parameters
      'product_agg': product_agg,
      'agg_true_book': 10000.,
      'agg_lot_size_min': 100,  # USD
      'agg_lot_size_max': 1.5 * lot_size,  # USD
      'agg_edge_bp': -5.,
      'agg_pos_multiplier': 100.
  }

  # Trading environment config
  if init_og_immediately:
    init_og_after_min = 0
  else:
    init_og_after_min = strategy_params['basis_ma_window_min']

  exit_after_min = init_og_after_min + 50
  env_config = {
      'first_product': product_agg,
      'first_og_config_name': okex_futures_og_config,
      'second_product': product_passive,
      'second_og_config_name': bitmex_og_config,
      'init_order_gateway_after': datetime.timedelta(minutes=init_og_after_min),
      'exit_after': datetime.timedelta(minutes=exit_after_min)
  }

  return strategy_params, env_config


def main(argv):
  if len(argv) == 2:
    param_set_name = argv[1]
  else:
    param_set_name = 'quarter'

  logging.info('Parameter set name: %s', param_set_name)
  strategy_params, env_config = get_param_set(param_set_name, False)

  # Run strategy!
  strategy = SoftArbStrategy(strategy_params, dry_run=False)
  return StrategyBase.run_live(on_base_init_callback=strategy.on_base_init,
                               on_book_reset_callback=strategy.on_book_reset,
                               on_order_gateway_reset_callback=strategy.on_og_reset,
                               **env_config)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
